import json
import csv
import yaml
import pandas as pd
from pathlib import Path
from typing import Dict, Any, Optional, Union, List
from .logger import test_logger

class TestDataManager:
    """测试数据管理器"""
    
    def __init__(self, data_dir: str = "test_data"):
        self.data_dir = Path(__file__).parent.parent / data_dir
        self.data_dir.mkdir(exist_ok=True)
        self.logger = test_logger.get_logger(self.__class__.__name__)
        self._current_data: Dict[str, Any] = {}
        
    def load_data(self, module: str, data_type: str, format: str = 'json') -> Dict[str, Any]:
        """加载测试数据
        
        Args:
            module: 模块名称
            data_type: 数据类型
            format: 数据格式（json/csv/excel/yaml）
            
        Returns:
            测试数据字典
        """
        file_path = self._get_data_path(module, data_type, format)
        if not file_path.exists():
            return {}
            
        try:
            if format.lower() == 'json':
                return self._load_json(file_path)
            elif format.lower() == 'csv':
                return self._load_csv(file_path)
            elif format.lower() == 'excel':
                return self._load_excel(file_path)
            elif format.lower() == 'yaml':
                return self._load_yaml(file_path)
            else:
                raise ValueError(f"Unsupported data format: {format}")
        except Exception as e:
            self.logger.error(f"Failed to load data from {file_path}: {str(e)}")
            return {}
            
    def save_data(self, module: str, data_type: str, data: Dict[str, Any], format: str = 'json'):
        """保存测试数据
        
        Args:
            module: 模块名称
            data_type: 数据类型
            data: 测试数据
            format: 数据格式（json/csv/excel/yaml）
        """
        file_path = self._get_data_path(module, data_type, format)
        
        try:
            if format.lower() == 'json':
                self._save_json(file_path, data)
            elif format.lower() == 'csv':
                self._save_csv(file_path, data)
            elif format.lower() == 'excel':
                self._save_excel(file_path, data)
            elif format.lower() == 'yaml':
                self._save_yaml(file_path, data)
            else:
                raise ValueError(f"Unsupported data format: {format}")
        except Exception as e:
            self.logger.error(f"Failed to save data to {file_path}: {str(e)}")
            
    def get_data(self, key: str) -> Any:
        """获取当前测试数据"""
        return self._current_data.get(key)
        
    def set_data(self, key: str, value: Any):
        """设置当前测试数据"""
        self._current_data[key] = value
        self.logger.info(f"Set test data: {key} = {value}")
        
    def clear_data(self):
        """清除当前测试数据"""
        self._current_data.clear()
        
    def _get_data_path(self, module: str, data_type: str, format: str) -> Path:
        """获取数据文件路径"""
        return self.data_dir / f"{module}_{data_type}.{format.lower()}"
        
    def _load_json(self, file_path: Path) -> Dict[str, Any]:
        """加载JSON数据"""
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
            
    def _save_json(self, file_path: Path, data: Dict[str, Any]):
        """保存JSON数据"""
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
            
    def _load_csv(self, file_path: Path) -> Dict[str, Any]:
        """加载CSV数据"""
        data = {}
        with open(file_path, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            for row in reader:
                key = row.get('key')
                if key:
                    data[key] = row
        return data
        
    def _save_csv(self, file_path: Path, data: Dict[str, Any]):
        """保存CSV数据"""
        # 获取所有字段名
        fieldnames = set()
        for value in data.values():
            fieldnames.update(value.keys())
        fieldnames = ['key'] + sorted(list(fieldnames - {'key'}))
        
        # 写入CSV文件
        with open(file_path, 'w', encoding='utf-8', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            for key, value in data.items():
                row = {'key': key, **value}
                writer.writerow(row)
                
    def _load_excel(self, file_path: Path) -> Dict[str, Any]:
        """加载Excel数据"""
        df = pd.read_excel(file_path)
        data = {}
        for _, row in df.iterrows():
            key = row.get('key')
            if key:
                data[key] = row.to_dict()
        return data
        
    def _save_excel(self, file_path: Path, data: Dict[str, Any]):
        """保存Excel数据"""
        # 将字典转换为DataFrame
        rows = []
        for key, value in data.items():
            row = {'key': key, **value}
            rows.append(row)
            
        df = pd.DataFrame(rows)
        df.to_excel(file_path, index=False)
        
    def _load_yaml(self, file_path: Path) -> Dict[str, Any]:
        """加载YAML数据"""
        with open(file_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
            
    def _save_yaml(self, file_path: Path, data: Dict[str, Any]):
        """保存YAML数据"""
        with open(file_path, 'w', encoding='utf-8') as f:
            yaml.safe_dump(data, f, allow_unicode=True, sort_keys=False)
            
    def load_test_cases(self, module: str) -> List[Dict[str, Any]]:
        """加载测试用例
        
        Args:
            module: 模块名称
            
        Returns:
            测试用例列表
        """
        file_path = self.data_dir / f"{module}_test_cases.yaml"
        if not file_path.exists():
            return []
            
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                test_cases = yaml.safe_load(f)
                return test_cases if isinstance(test_cases, list) else []
        except Exception as e:
            self.logger.error(f"Failed to load test cases from {file_path}: {str(e)}")
            return []
            
    def save_test_cases(self, module: str, test_cases: List[Dict[str, Any]]):
        """保存测试用例
        
        Args:
            module: 模块名称
            test_cases: 测试用例列表
        """
        file_path = self.data_dir / f"{module}_test_cases.yaml"
        
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                yaml.safe_dump(test_cases, f, allow_unicode=True, sort_keys=False)
        except Exception as e:
            self.logger.error(f"Failed to save test cases to {file_path}: {str(e)}") 